26class ArgPack:
 27    """ The `ArgPack` Type Class.
 28
 29    The `ArgPack` operates as a dictionary-like data pack, storing members as (key, value) pairs. Members stored can
 30    range from scalars and matrices to other dictionary-like structures. Distinguished from structs, `ArgPack` can
 31    accommodate buffer types such as `NdarrayType` and `TextureType` from Taichi. However, unlike `ti.Struct` which
 32    serves as a data container, `ArgPack` functions as a reference container. It's important to note that `ArgPack`
 33    cannot be nested within other types except for another `ArgPack`, and can only be utilized as kernel parameters.
 34
 35    Args:
 36        annotations (Dict[str, Union[Dict, Matrix, Struct]]): \
 37            The keys and types for `ArgPack` members.
 38        dtype (ArgPackType): \
 39            The ArgPackType class of this ArgPack object.
 40        entries (Dict[str, Union[Dict, Matrix, Struct]]): \
 41            The keys and corresponding values for `ArgPack` members.
 42
 43    Returns:
 44        An instance of this `ArgPack`.
 45
 46    Example::
 47
 48        >>> vec3 = ti.types.vector(3, ti.f32)
 49        >>> pack_type = ti.ArgPackType(v=vec3, t=ti.f32)
 50        >>> a = pack_type(v=vec3([0, 0, 0]), t=1.0)
 51        >>> print(a.items)
 52        dict_items([('v', [0. 0. 0.]), ('t', 1.0)])
 53    """
 54
 55    _instance_count = 0
 56
 57    def __init__(self, annotations, dtype, *args, **kwargs):
 58        # converts dicts to argument packs
 59        if len(args) == 1 and kwargs == {} and isinstance(args[0], dict):
 60            self.__entries = args[0]
 61        elif len(args) == 0:
 62            self.__entries = kwargs
 63        else:
 64            raise TaichiSyntaxError(
 65                "Custom argument packs need to be initialized using either dictionary or keyword arguments"
 66            )
 67        if annotations.keys() != self.__entries.keys():
 68            raise TaichiSyntaxError("ArgPack annotations keys not equals to entries keys.")
 69        self.__annotations = annotations
 70        for k, v in self.__entries.items():
 71            self.__entries[k] = v if in_python_scope() else impl.expr_init(v)
 72        self._register_members()
 73        self.__dtype = dtype
 74        self.__argpack = impl.get_runtime().prog.create_argpack(self.__dtype)
 75        for i, (k, v) in enumerate(self.__entries.items()):
 76            self._write_to_device(self.__annotations[k], type(v), v, self._calc_element_true_index(i))
 77
 78    def __del__(self):
 79        if impl is not None and impl.get_runtime() is not None and impl.get_runtime().prog is not None:
 80            impl.get_runtime().prog.delete_argpack(self.__argpack)
 81
 82    @property
 83    def keys(self):
 84        """Returns the list of member names in string format.
 85
 86        Example::
 87
 88           >>> vec3 = ti.types.vector(3, ti.f32)
 89           >>> sphere_pack = ti.ArgPackType(center=vec3, radius=ti.f32)
 90           >>> sphere = sphere_pack(center=vec3([0, 0, 0]), radius=1.0)
 91           >>> sphere.keys
 92           ['center', 'radius']
 93        """
 94        return list(self.__entries.keys())
 95
 96    @property
 97    def _members(self):
 98        return list(self.__entries.values())
 99
100    @property
101    def _annotations(self):
102        return list(self.__annotations.values())
103
104    @property
105    def items(self):
106        """Returns the items in this argument pack.
107
108        Example::
109
110           >>> vec3 = ti.types.vector(3, ti.f32)
111           >>> sphere_pack = ti.ArgPackType(center=vec3, radius=ti.f32)
112           >>> sphere = sphere_pack(center=vec3([0, 0, 0]), radius=1.0)
113           >>> sphere.items
114            dict_items([('center', 2), ('radius', 1.0)])
115        """
116        return self.__entries.items()
117
118    def __getitem__(self, key):
119        ret = self.__entries[key]
120        return ret
121
122    def __setitem__(self, key, value):
123        self.__entries[key] = value
124        index = self._calc_element_true_index(list(self.__annotations).index(key))
125        self._write_to_device(self.__annotations[key], type(value), value, index)
126
127    def _set_entries(self, value):
128        if isinstance(value, dict):
129            value = ArgPack(self.__annotations, value)
130        for k in self.keys:
131            self[k] = value[k]
132
133    @staticmethod
134    def _make_getter(key):
135        def getter(self):
136            """Get an entry from custom argument pack by name."""
137            return self[key]
138
139        return getter
140
141    @staticmethod
142    def _make_setter(key):
143        @python_scope
144        def setter(self, value):
145            self[key] = value
146
147        return setter
148
149    def _register_members(self):
150        # https://stackoverflow.com/questions/48448074/adding-a-property-to-an-existing-object-instance
151        cls = self.__class__
152        new_cls_name = cls.__name__ + str(cls._instance_count)
153        cls._instance_count += 1
154        properties = {k: property(cls._make_getter(k), cls._make_setter(k)) for k in self.keys}
155        self.__class__ = type(new_cls_name, (cls,), properties)
156
157    def __len__(self):
158        """Get the number of entries in a custom argument pack."""
159        return len(self.__entries)
160
161    def __iter__(self):
162        return self.__entries.values()
163
164    def __str__(self):
165        """Python scope argument pack array print support."""
166        if impl.inside_kernel():
167            item_str = ", ".join([str(k) + "=" + str(v) for k, v in self.items])
168            return f"<ti.ArgPack {item_str}>"
169        return str(self.to_dict())
170
171    def __repr__(self):
172        return str(self.to_dict())
173
174    def to_dict(self):
175        """Converts the ArgPack to a dictionary.
176
177        Returns:
178            Dict: The result dictionary.
179        """
180        res_dict = {
181            k: v.to_dict() if isinstance(v, ArgPack) else v.to_list() if isinstance(v, Matrix) else v
182            for k, v in self.__entries.items()
183        }
184        return res_dict
185
186    def _calc_element_true_index(self, old_index):
187        for i in range(old_index):
188            anno = list(self.__annotations.values())[i]
189            if (
190                isinstance(anno, sparse_matrix_builder)
191                or isinstance(anno, ndarray_type.NdarrayType)
192                or isinstance(anno, texture_type.TextureType)
193                or isinstance(anno, texture_type.RWTextureType)
194                or isinstance(anno, ndarray_type.NdarrayType)
195            ):
196                old_index -= 1
197        return old_index
198
199    def _write_to_device(self, needed, provided, v, index):
200        if isinstance(needed, ArgPackType):
201            if not isinstance(v, ArgPack):
202                raise TaichiRuntimeTypeError.get(index, str(needed), str(provided))
203            self.__argpack.set_arg_nested_argpack(index, v.__argpack)
204        else:
205            # Note: do not use sth like "needed == f32". That would be slow.
206            if id(needed) in primitive_types.real_type_ids:
207                if not isinstance(v, (float, int, np.floating, np.integer)):
208                    raise TaichiRuntimeTypeError.get(index, needed.to_string(), provided)
209                self.__argpack.set_arg_float((index,), float(v))
210            elif id(needed) in primitive_types.integer_type_ids:
211                if not isinstance(v, (int, np.integer)):
212                    raise TaichiRuntimeTypeError.get(index, needed.to_string(), provided)
213                if is_signed(cook_dtype(needed)):
214                    self.__argpack.set_arg_int((index,), int(v))
215                else:
216                    self.__argpack.set_arg_uint((index,), int(v))
217            elif isinstance(needed, sparse_matrix_builder):
218                pass
219            elif isinstance(needed, ndarray_type.NdarrayType) and isinstance(v, taichi.lang._ndarray.Ndarray):
220                pass
221            elif isinstance(needed, texture_type.TextureType) and isinstance(v, taichi.lang._texture.Texture):
222                pass
223            elif isinstance(needed, texture_type.RWTextureType) and isinstance(v, taichi.lang._texture.Texture):
224                pass
225            elif isinstance(needed, ndarray_type.NdarrayType):
226                pass
227            elif isinstance(needed, MatrixType):
228                if needed.dtype in primitive_types.real_types:
229
230                    def cast_func(x):
231                        if not isinstance(x, (int, float, np.integer, np.floating)):
232                            raise TaichiRuntimeTypeError.get(index, needed.dtype.to_string(), type(x))
233                        return float(x)
234
235                elif needed.dtype in primitive_types.integer_types:
236
237                    def cast_func(x):
238                        if not isinstance(x, (int, np.integer)):
239                            raise TaichiRuntimeTypeError.get(index, needed.dtype.to_string(), type(x))
240                        return int(x)
241
242                else:
243                    raise ValueError(f"Matrix dtype {needed.dtype} is not integer type or real type.")
244
245                if needed.ndim == 2:
246                    v = [cast_func(v[i, j]) for i in range(needed.n) for j in range(needed.m)]
247                else:
248                    v = [cast_func(v[i]) for i in range(needed.n)]
249                v = needed(*v)
250                needed.set_argpack_struct_args(v, self.__argpack, (index,))
251            elif isinstance(needed, StructType):
252                if not isinstance(v, needed):
253                    raise TaichiRuntimeTypeError.get(index, str(needed), provided)
254                needed.set_argpack_struct_args(v, self.__argpack, (index,))
255            else:
256                raise ValueError(f"Argument type mismatch. Expecting {needed}, got {type(v)}.")

433class BitpackedFields:
434    """Taichi bitpacked fields, where fields with quantized types are packed together.
435
436    Args:
437        max_num_bits (int): Maximum number of bits all fields inside can occupy in total. Only 32 or 64 is allowed.
438    """
439
440    def __init__(self, max_num_bits):
441        self.fields = []
442        self.bit_struct_type_builder = _ti_core.BitStructTypeBuilder(max_num_bits)
443
444    def place(self, *args, shared_exponent=False):
445        """Places a list of fields with quantized types inside.
446
447        Args:
448            *args (List[Field]): A list of fields with quantized types to place.
449            shared_exponent (bool): Whether the fields have a shared exponent.
450        """
451        if shared_exponent:
452            self.bit_struct_type_builder.begin_placing_shared_exponent()
453        count = 0
454        for arg in args:
455            assert isinstance(arg, Field)
456            for var in arg._get_field_members():
457                self.fields.append((var.ptr, self.bit_struct_type_builder.add_member(var.ptr.get_dt())))
458                count += 1
459        if shared_exponent:
460            self.bit_struct_type_builder.end_placing_shared_exponent()
461            if count <= 1:
462                raise TaichiSyntaxError("At least 2 fields need to be placed when shared_exponent=True")

21class DeviceCapability:
22    spirv_version_1_3 = "spirv_version=66304"
23    spirv_version_1_4 = "spirv_version=66560"
24    spirv_version_1_5 = "spirv_version=66816"
25    spirv_has_int8 = "spirv_has_int8"
26    spirv_has_int16 = "spirv_has_int16"
27    spirv_has_int64 = "spirv_has_int64"
28    spirv_has_float16 = "spirv_has_float16"
29    spirv_has_float64 = "spirv_has_float64"
30    spirv_has_atomic_int64 = "spirv_has_atomic_int64"
31    spirv_has_atomic_float16 = "spirv_has_atomic_float16"
32    spirv_has_atomic_float16_add = "spirv_has_atomic_float16_add"
33    spirv_has_atomic_float16_minmax = "spirv_has_atomic_float16_minmax"
34    spirv_has_atomic_float = "spirv_has_atomic_float"
35    spirv_has_atomic_float_add = "spirv_has_atomic_float_add"
36    spirv_has_atomic_float_minmax = "spirv_has_atomic_float_minmax"
37    spirv_has_atomic_float64 = "spirv_has_atomic_float64"
38    spirv_has_atomic_float64_add = "spirv_has_atomic_float64_add"
39    spirv_has_atomic_float64_minmax = "spirv_has_atomic_float64_minmax"
40    spirv_has_variable_ptr = "spirv_has_variable_ptr"
41    spirv_has_physical_storage_buffer = "spirv_has_physical_storage_buffer"
42    spirv_has_subgroup_basic = "spirv_has_subgroup_basic"
43    spirv_has_subgroup_vote = "spirv_has_subgroup_vote"
44    spirv_has_subgroup_arithmetic = "spirv_has_subgroup_arithmetic"
45    spirv_has_subgroup_ballot = "spirv_has_subgroup_ballot"
46    spirv_has_non_semantic_info = "spirv_has_non_semantic_info"
47    spirv_has_no_integer_wrap_decoration = "spirv_has_no_integer_wrap_decoration"

 18class Field:
 19    """Taichi field class.
 20
 21    A field is constructed by a list of field members.
 22    For example, a scalar field has 1 field member, while a 3x3 matrix field has 9 field members.
 23    A field member is a Python Expr wrapping a C++ FieldExpression.
 24
 25    Args:
 26        vars (List[Expr]): Field members.
 27    """
 28
 29    def __init__(self, _vars):
 30        assert all(_vars)
 31        self.vars = _vars
 32        self.host_accessors = None
 33        self.grad = None
 34        self.dual = None
 35
 36    @property
 37    def snode(self):
 38        """Gets representative SNode for info purposes.
 39
 40        Returns:
 41            SNode: Representative SNode (SNode of first field member).
 42        """
 43        return self._snode
 44
 45    @property
 46    def _snode(self):
 47        """Gets representative SNode for info purposes.
 48
 49        Returns:
 50            SNode: Representative SNode (SNode of first field member).
 51        """
 52        return taichi.lang.snode.SNode(self.vars[0].ptr.snode())
 53
 54    @property
 55    def shape(self):
 56        """Gets field shape.
 57
 58        Returns:
 59            Tuple[Int]: Field shape.
 60        """
 61        return self._snode.shape
 62
 63    @property
 64    def dtype(self):
 65        """Gets data type of each individual value.
 66
 67        Returns:
 68            DataType: Data type of each individual value.
 69        """
 70        return self._snode._dtype
 71
 72    @property
 73    def _name(self):
 74        """Gets field name.
 75
 76        Returns:
 77            str: Field name.
 78        """
 79        return self._snode._name
 80
 81    def parent(self, n=1):
 82        """Gets an ancestor of the representative SNode in the SNode tree.
 83
 84        Args:
 85            n (int): the number of levels going up from the representative SNode.
 86
 87        Returns:
 88            SNode: The n-th parent of the representative SNode.
 89        """
 90        return self.snode.parent(n)
 91
 92    def _get_field_members(self):
 93        """Gets field members.
 94
 95        Returns:
 96            List[Expr]: Field members.
 97        """
 98        return self.vars
 99
100    def _loop_range(self):
101        """Gets SNode of representative field member for loop range info.
102
103        Returns:
104            taichi_python.SNode: SNode of representative (first) field member.
105        """
106        return self.vars[0].ptr.snode()
107
108    def _set_grad(self, grad):
109        """Sets corresponding grad field (reverse mode).
110        Args:
111            grad (Field): Corresponding grad field.
112        """
113        self.grad = grad
114
115    def _set_dual(self, dual):
116        """Sets corresponding dual field (forward mode).
117
118        Args:
119            dual (Field): Corresponding dual field.
120        """
121        self.dual = dual
122
123    @python_scope
124    def fill(self, val):
125        """Fills `self` with a specific value.
126
127        Args:
128            val (Union[int, float]): Value to fill.
129        """
130        raise NotImplementedError()
131
132    @python_scope
133    def to_numpy(self, dtype=None):
134        """Converts `self` to a numpy array.
135
136        Args:
137            dtype (DataType, optional): The desired data type of returned numpy array.
138
139        Returns:
140            numpy.ndarray: The result numpy array.
141        """
142        raise NotImplementedError()
143
144    @python_scope
145    def to_torch(self, device=None):
146        """Converts `self` to a torch tensor.
147
148        Args:
149            device (torch.device, optional): The desired device of returned tensor.
150
151        Returns:
152            torch.tensor: The result torch tensor.
153        """
154        raise NotImplementedError()
155
156    @python_scope
157    def to_paddle(self, place=None):
158        """Converts `self` to a paddle tensor.
159
160        Args:
161            place (paddle.CPUPlace()/CUDAPlace(n), optional): The desired place of returned tensor.
162
163        Returns:
164            paddle.Tensor: The result paddle tensor.
165        """
166        raise NotImplementedError()
167
168    @python_scope
169    def from_numpy(self, arr):
170        """Loads all elements from a numpy array.
171
172        The shape of the numpy array needs to be the same as `self`.
173
174        Args:
175            arr (numpy.ndarray): The source numpy array.
176        """
177        raise NotImplementedError()
178
179    @python_scope
180    def _from_external_arr(self, arr):
181        raise NotImplementedError()
182
183    @python_scope
184    def from_torch(self, arr):
185        """Loads all elements from a torch tensor.
186
187        The shape of the torch tensor needs to be the same as `self`.
188
189        Args:
190            arr (torch.tensor): The source torch tensor.
191        """
192        self._from_external_arr(arr.contiguous())
193
194    @python_scope
195    def from_paddle(self, arr):
196        """Loads all elements from a paddle tensor.
197
198        The shape of the paddle tensor needs to be the same as `self`.
199
200        Args:
201            arr (paddle.Tensor): The source paddle tensor.
202        """
203        self.from_numpy(arr)
204
205    @python_scope
206    def copy_from(self, other):
207        """Copies all elements from another field.
208
209        The shape of the other field needs to be the same as `self`.
210
211        Args:
212            other (Field): The source field.
213        """
214        if not isinstance(other, Field):
215            raise TypeError("Cannot copy from a non-field object")
216        if self.shape != other.shape:
217            raise ValueError(f"ti.field shape {self.shape} does not match" f" the source field shape {other.shape}")
218        from taichi._kernels import tensor_to_tensor  # pylint: disable=C0415
219
220        tensor_to_tensor(self, other)
221
222    @python_scope
223    def __setitem__(self, key, value):
224        """Sets field element in Python scope.
225
226        Args:
227            key (Union[List[int], int, None]): Coordinates of the field element.
228            value (element type): Value to set.
229        """
230        raise NotImplementedError()
231
232    @python_scope
233    def __getitem__(self, key):
234        """Gets field element in Python scope.
235
236        Args:
237            key (Union[List[int], int, None]): Coordinates of the field element.
238
239        Returns:
240            element type: Value retrieved.
241        """
242        raise NotImplementedError()
243
244    def __str__(self):
245        if taichi.lang.impl.inside_kernel():
246            return self.__repr__()  # make pybind11 happy, see Matrix.__str__
247        if self._snode.ptr is None:
248            return "<Field: Definition of this field is incomplete>"
249        return str(self.to_numpy())
250
251    def _pad_key(self, key):
252        if key is None:
253            key = ()
254        if not isinstance(key, (tuple, list)):
255            key = (key,)
256
257        if len(key) != len(self.shape):
258            raise AssertionError("Slicing is not supported on ti.field")
259
260        return key + ((0,) * (_ti_core.get_max_num_indices() - len(key)))
261
262    def _initialize_host_accessors(self):
263        if self.host_accessors:
264            return
265        taichi.lang.impl.get_runtime().materialize()
266        self.host_accessors = [SNodeHostAccessor(e.ptr.snode()) for e in self.vars]
267
268    def _host_access(self, key):
269        return [SNodeHostAccess(e, key) for e in self.host_accessors]
270
271    def __iter__(self):
272        raise NotImplementedError("Struct for is only available in Taichi scope.")

 207@_gen_swizzles
 208class Matrix(TaichiOperations):
 209    """The matrix class.
 210
 211    A matrix is a 2-D rectangular array with scalar entries, it's row-majored, and is
 212    aligned continuously. We recommend only use matrix with no more than 32 elements for
 213    efficiency considerations.
 214
 215    Note: in taichi a matrix is strictly two-dimensional and only stores scalars.
 216
 217    Args:
 218        arr (Union[list, tuple, np.ndarray]): the initial values of a matrix.
 219        dt (:mod:`~taichi.types.primitive_types`): the element data type.
 220        ndim (int optional): the number of dimensions of the matrix; forced reshape if given.
 221
 222    Example::
 223
 224        use a 2d list to initialize a matrix
 225
 226        >>> @ti.kernel
 227        >>> def test():
 228        >>>     n = 5
 229        >>>     M = ti.Matrix([[0] * n for _ in range(n)], ti.i32)
 230        >>>     print(M)  # a 5x5 matrix with integer elements
 231
 232        get the number of rows and columns via the `n`, `m` property:
 233
 234        >>> M = ti.Matrix([[0, 1], [2, 3], [4, 5]], ti.i32)
 235        >>> M.n  # number of rows
 236        3
 237        >>> M.m  # number of cols
 238        >>> 2
 239
 240        you can even initialize a matrix with an empty list:
 241
 242        >>> M = ti.Matrix([[], []], ti.i32)
 243        >>> M.n
 244        2
 245        >>> M.m
 246        0
 247    """
 248
 249    _is_taichi_class = True
 250    _is_matrix_class = True
 251    __array_priority__ = 1000
 252
 253    def __init__(self, arr, dt=None):
 254        if not isinstance(arr, (list, tuple, np.ndarray)):
 255            raise TaichiTypeError("An Matrix/Vector can only be initialized with an array-like object")
 256        if len(arr) == 0:
 257            self.ndim = 0
 258            self.n, self.m = 0, 0
 259            self.entries = np.array([])
 260            self.is_host_access = False
 261        elif isinstance(arr[0], Matrix):
 262            raise Exception("cols/rows required when using list of vectors")
 263        elif isinstance(arr[0], Iterable):  # matrix
 264            self.ndim = 2
 265            self.n, self.m = len(arr), len(arr[0])
 266            if isinstance(arr[0][0], (SNodeHostAccess, NdarrayHostAccess)):
 267                self.entries = arr
 268                self.is_host_access = True
 269            else:
 270                self.entries = np.array(arr, None if dt is None else to_numpy_type(dt))
 271                self.is_host_access = False
 272        else:  # vector
 273            self.ndim = 1
 274            self.n, self.m = len(arr), 1
 275            if isinstance(arr[0], (SNodeHostAccess, NdarrayHostAccess)):
 276                self.entries = arr
 277                self.is_host_access = True
 278            else:
 279                self.entries = np.array(arr, None if dt is None else to_numpy_type(dt))
 280                self.is_host_access = False
 281
 282        if self.n * self.m > 32:
 283            warning(
 284                f"Taichi matrices/vectors with {self.n}x{self.m} > 32 entries are not suggested."
 285                " Matrices/vectors will be automatically unrolled at compile-time for performance."
 286                " So the compilation time could be extremely long if the matrix size is too big."
 287                " You may use a field to store a large matrix like this, e.g.:\n"
 288                f"    x = ti.field(ti.f32, ({self.n}, {self.m})).\n"
 289                " See https://docs.taichi-lang.org/docs/field#matrix-size"
 290                " for more details.",
 291                UserWarning,
 292                stacklevel=2,
 293            )
 294
 295    def get_shape(self):
 296        if self.ndim == 1:
 297            return (self.n,)
 298        if self.ndim == 2:
 299            return (self.n, self.m)
 300        return None
 301
 302    def __matmul__(self, other):
 303        """Matrix-matrix or matrix-vector multiply.
 304
 305        Args:
 306            other (Union[Matrix, Vector]): a matrix or a vector.
 307
 308        Returns:
 309            The matrix-matrix product or matrix-vector product.
 310
 311        """
 312        from taichi.lang import matrix_ops  # pylint: disable=C0415
 313
 314        return matrix_ops.matmul(self, other)
 315
 316    # host access & python scope operation
 317    def __len__(self):
 318        """Get the length of each row of a matrix"""
 319        # TODO: When this is a vector, should return its dimension?
 320        return self.n
 321
 322    def __iter__(self):
 323        if self.ndim == 1:
 324            return (self[i] for i in range(self.n))
 325        return ([self[i, j] for j in range(self.m)] for i in range(self.n))
 326
 327    def __getitem__(self, indices):
 328        """Access to the element at the given indices in a matrix.
 329
 330        Args:
 331            indices (Sequence[Expr]): the indices of the element.
 332
 333        Returns:
 334            The value of the element at a specific position of a matrix.
 335
 336        """
 337        entry = self._get_entry(indices)
 338        if self.is_host_access:
 339            return _read_host_access(entry)
 340        return entry
 341
 342    @python_scope
 343    def __setitem__(self, indices, item):
 344        """Set the element value at the given indices in a matrix.
 345
 346        Args:
 347            indices (Sequence[Expr]): the indices of a element.
 348
 349        """
 350        if self.is_host_access:
 351            entry = self._get_entry(indices)
 352            _write_host_access(entry, item)
 353        else:
 354            if not isinstance(indices, (list, tuple)):
 355                indices = [indices]
 356            assert len(indices) in [1, 2]
 357            assert len(indices) == self.ndim, f"Expected {self.ndim} indices, got {len(indices)}"
 358            if self.ndim == 1:
 359                self.entries[indices[0]] = item
 360            else:
 361                self.entries[indices[0]][indices[1]] = item
 362
 363    def _get_entry(self, indices):
 364        if not isinstance(indices, (list, tuple)):
 365            indices = [indices]
 366        assert len(indices) in [1, 2]
 367        assert len(indices) == self.ndim, f"Expected {self.ndim} indices, got {len(indices)}"
 368        if self.ndim == 1:
 369            return self.entries[indices[0]]
 370        return self.entries[indices[0]][indices[1]]
 371
 372    def _get_slice(self, a, b):
 373        if isinstance(a, slice):
 374            a = range(a.start or 0, a.stop or self.n, a.step or 1)
 375        if isinstance(b, slice):
 376            b = range(b.start or 0, b.stop or self.m, b.step or 1)
 377        if isinstance(a, range) and isinstance(b, range):
 378            return Matrix([[self._get_entry(i, j) for j in b] for i in a])
 379        if isinstance(a, range):  # b is not range
 380            return Vector([self._get_entry(i, b) for i in a])
 381        # a is not range while b is range
 382        return Vector([self._get_entry(a, j) for j in b])
 383
 384    @python_scope
 385    def _set_entries(self, value):
 386        if isinstance(value, Matrix):
 387            value = value.to_list()
 388        if self.is_host_access:
 389            if self.ndim == 1:
 390                for i in range(self.n):
 391                    _write_host_access(self.entries[i], value[i])
 392            else:
 393                for i in range(self.n):
 394                    for j in range(self.m):
 395                        _write_host_access(self.entries[i][j], value[i][j])
 396        else:
 397            if self.ndim == 1:
 398                for i in range(self.n):
 399                    self.entries[i] = value[i]
 400            else:
 401                for i in range(self.n):
 402                    for j in range(self.m):
 403                        self.entries[i][j] = value[i][j]
 404
 405    @property
 406    def _members(self):
 407        return self.entries
 408
 409    def to_list(self):
 410        """Return this matrix as a 1D `list`.
 411
 412        This is similar to `numpy.ndarray`'s `flatten` and `ravel` methods,
 413        the difference is that this function always returns a new list.
 414        """
 415        if self.is_host_access:
 416            if self.ndim == 1:
 417                return [_read_host_access(self.entries[i]) for i in range(self.n)]
 418            assert self.ndim == 2
 419            return [[_read_host_access(self.entries[i][j]) for j in range(self.m)] for i in range(self.n)]
 420        return self.entries.tolist()
 421
 422    @taichi_scope
 423    def cast(self, dtype):
 424        """Cast the matrix elements to a specified data type.
 425
 426        Args:
 427            dtype (:mod:`~taichi.types.primitive_types`): data type of the
 428                returned matrix.
 429
 430        Returns:
 431            :class:`taichi.Matrix`: A new matrix with the specified data dtype.
 432
 433        Example::
 434
 435            >>> A = ti.Matrix([0, 1, 2], ti.i32)
 436            >>> B = A.cast(ti.f32)
 437            >>> B
 438            [0.0, 1.0, 2.0]
 439        """
 440        if self.ndim == 1:
 441            return Vector([ops_mod.cast(self[i], dtype) for i in range(self.n)])
 442        return Matrix([[ops_mod.cast(self[i, j], dtype) for j in range(self.m)] for i in range(self.n)])
 443
 444    def trace(self):
 445        """The sum of a matrix diagonal elements.
 446
 447        To call this method the matrix must be square-like.
 448
 449        Returns:
 450            The sum of a matrix diagonal elements.
 451
 452        Example::
 453
 454            >>> m = ti.Matrix([[1, 2], [3, 4]])
 455            >>> m.trace()
 456            5
 457        """
 458        # pylint: disable-msg=C0415
 459        from taichi.lang import matrix_ops
 460
 461        return matrix_ops.trace(self)
 462
 463    def inverse(self):
 464        """Returns the inverse of this matrix.
 465
 466        Note:
 467            The matrix dimension should be less than or equal to 4.
 468
 469        Returns:
 470            :class:`~taichi.Matrix`: The inverse of a matrix.
 471
 472        Raises:
 473            Exception: Inversions of matrices with sizes >= 5 are not supported.
 474        """
 475        from taichi.lang import matrix_ops  # pylint: disable=C0415
 476
 477        return matrix_ops.inverse(self)
 478
 479    def normalized(self, eps=0):
 480        """Normalize a vector, i.e. matrices with the second dimension being
 481        equal to one.
 482
 483        The normalization of a vector `v` is a vector of length 1
 484        and has the same direction with `v`. It's equal to `v/|v|`.
 485
 486        Args:
 487            eps (float): a safe-guard value for sqrt, usually 0.
 488
 489        Example::
 490
 491            >>> a = ti.Vector([3, 4], ti.f32)
 492            >>> a.normalized()
 493            [0.6, 0.8]
 494        """
 495        # pylint: disable-msg=C0415
 496        from taichi.lang import matrix_ops
 497
 498        return matrix_ops.normalized(self, eps)
 499
 500    def transpose(self):
 501        """Returns the transpose of a matrix.
 502
 503        Returns:
 504            :class:`~taichi.Matrix`: The transpose of this matrix.
 505
 506        Example::
 507
 508            >>> A = ti.Matrix([[0, 1], [2, 3]])
 509            >>> A.transpose()
 510            [[0, 2], [1, 3]]
 511        """
 512        # pylint: disable=C0415
 513        from taichi.lang import matrix_ops
 514
 515        return matrix_ops.transpose(self)
 516
 517    @taichi_scope
 518    def determinant(a):
 519        """Returns the determinant of this matrix.
 520
 521        Note:
 522            The matrix dimension should be less than or equal to 4.
 523
 524        Returns:
 525            dtype: The determinant of this matrix.
 526
 527        Raises:
 528            Exception: Determinants of matrices with sizes >= 5 are not supported.
 529        """
 530        # pylint: disable=C0415
 531        from taichi.lang import matrix_ops
 532
 533        return matrix_ops.determinant(a)
 534
 535    @staticmethod
 536    def diag(dim, val):
 537        """Returns a diagonal square matrix with the diagonals filled
 538        with `val`.
 539
 540        Args:
 541            dim (int): the dimension of the wanted square matrix.
 542            val (TypeVar): value for the diagonal elements.
 543
 544        Returns:
 545            :class:`~taichi.Matrix`: The wanted diagonal matrix.
 546
 547        Example::
 548
 549            >>> m = ti.Matrix.diag(3, 1)
 550            [[1, 0, 0],
 551             [0, 1, 0],
 552             [0, 0, 1]]
 553        """
 554        # pylint: disable=C0415
 555        from taichi.lang import matrix_ops
 556
 557        return matrix_ops.diag(dim, val)
 558
 559    def sum(self):
 560        """Return the sum of all elements.
 561
 562        Example::
 563
 564            >>> m = ti.Matrix([[1, 2], [3, 4]])
 565            >>> m.sum()
 566            10
 567        """
 568        # pylint: disable=C0415
 569        from taichi.lang import matrix_ops
 570
 571        return matrix_ops.sum(self)
 572
 573    def norm(self, eps=0):
 574        """Returns the square root of the sum of the absolute squares
 575        of its elements.
 576
 577        Args:
 578            eps (Number): a safe-guard value for sqrt, usually 0.
 579
 580        Example::
 581
 582            >>> a = ti.Vector([3, 4])
 583            >>> a.norm()
 584            5
 585
 586        Returns:
 587            The square root of the sum of the absolute squares of its elements.
 588        """
 589        # pylint: disable=C0415
 590        from taichi.lang import matrix_ops
 591
 592        return matrix_ops.norm(self, eps=eps)
 593
 594    def norm_inv(self, eps=0):
 595        """The inverse of the matrix :func:`~taichi.lang.matrix.Matrix.norm`.
 596
 597        Args:
 598            eps (float): a safe-guard value for sqrt, usually 0.
 599
 600        Returns:
 601            The inverse of the matrix/vector `norm`.
 602        """
 603        # pylint: disable=C0415
 604        from taichi.lang import matrix_ops
 605
 606        return matrix_ops.norm_inv(self, eps=eps)
 607
 608    def norm_sqr(self):
 609        """Returns the sum of the absolute squares of its elements."""
 610        # pylint: disable=C0415
 611        from taichi.lang import matrix_ops
 612
 613        return matrix_ops.norm_sqr(self)
 614
 615    def max(self):
 616        """Returns the maximum element value."""
 617        # pylint: disable=C0415
 618        from taichi.lang import matrix_ops
 619
 620        return matrix_ops.max(self)
 621
 622    def min(self):
 623        """Returns the minimum element value."""
 624        # pylint: disable=C0415
 625        from taichi.lang import matrix_ops
 626
 627        return matrix_ops.min(self)
 628
 629    def any(self):
 630        """Test whether any element not equal zero.
 631
 632        Returns:
 633            bool: `True` if any element is not equal zero, `False` otherwise.
 634
 635        Example::
 636
 637            >>> v = ti.Vector([0, 0, 1])
 638            >>> v.any()
 639            True
 640        """
 641        # pylint: disable=C0415
 642        from taichi.lang import matrix_ops
 643
 644        return matrix_ops.any(self)
 645
 646    def all(self):
 647        """Test whether all element not equal zero.
 648
 649        Returns:
 650            bool: `True` if all elements are not equal zero, `False` otherwise.
 651
 652        Example::
 653
 654            >>> v = ti.Vector([0, 0, 1])
 655            >>> v.all()
 656            False
 657        """
 658        # pylint: disable=C0415
 659        from taichi.lang import matrix_ops
 660
 661        return matrix_ops.all(self)
 662
 663    def fill(self, val):
 664        """Fills the matrix with a specified value.
 665
 666        Args:
 667            val (Union[int, float]): Value to fill.
 668
 669        Example::
 670
 671            >>> A = ti.Matrix([0, 1, 2, 3])
 672            >>> A.fill(-1)
 673            >>> A
 674            [-1, -1, -1, -1]
 675        """
 676        # pylint: disable=C0415
 677        from taichi.lang import matrix_ops
 678
 679        return matrix_ops.fill(self, val)
 680
 681    def to_numpy(self):
 682        """Converts this matrix to a numpy array.
 683
 684        Returns:
 685            numpy.ndarray: The result numpy array.
 686
 687        Example::
 688
 689            >>> A = ti.Matrix([[0], [1], [2], [3]])
 690            >>> A.to_numpy()
 691            >>> A
 692            array([[0], [1], [2], [3]])
 693        """
 694        if self.is_host_access:
 695            return np.array(self.to_list())
 696        return self.entries
 697
 698    @taichi_scope
 699    def __ti_repr__(self):
 700        yield "["
 701        for i in range(self.n):
 702            if i:
 703                yield ", "
 704            if self.m != 1:
 705                yield "["
 706            for j in range(self.m):
 707                if j:
 708                    yield ", "
 709                yield self(i, j)
 710            if self.m != 1:
 711                yield "]"
 712        yield "]"
 713
 714    def __str__(self):
 715        """Python scope matrix print support."""
 716        if impl.inside_kernel():
 717            """
 718            It seems that when pybind11 got an type mismatch, it will try
 719            to invoke `repr` to show the object... e.g.:
 720
 721            TypeError: make_const_expr_f32(): incompatible function arguments. The following argument types are supported:
 722                1. (arg0: float) -> taichi_python.Expr
 723
 724            Invoked with: <Taichi 2x1 Matrix>
 725
 726            So we have to make it happy with a dummy string...
 727            """
 728            return f"<{self.n}x{self.m} ti.Matrix>"
 729        return str(self.to_numpy())
 730
 731    def __repr__(self):
 732        return str(self.to_numpy())
 733
 734    @staticmethod
 735    @taichi_scope
 736    def zero(dt, n, m=None):
 737        """Constructs a Matrix filled with zeros.
 738
 739        Args:
 740            dt (DataType): The desired data type.
 741            n (int): The first dimension (row) of the matrix.
 742            m (int, optional): The second dimension (column) of the matrix.
 743
 744        Returns:
 745            :class:`~taichi.lang.matrix.Matrix`: A :class:`~taichi.lang.matrix.Matrix` instance filled with zeros.
 746
 747        """
 748        from taichi.lang import matrix_ops  # pylint: disable=C0415
 749
 750        if m is None:
 751            return matrix_ops._filled_vector(n, dt, 0)
 752        return matrix_ops._filled_matrix(n, m, dt, 0)
 753
 754    @staticmethod
 755    @taichi_scope
 756    def one(dt, n, m=None):
 757        """Constructs a Matrix filled with ones.
 758
 759        Args:
 760            dt (DataType): The desired data type.
 761            n (int): The first dimension (row) of the matrix.
 762            m (int, optional): The second dimension (column) of the matrix.
 763
 764        Returns:
 765            :class:`~taichi.lang.matrix.Matrix`: A :class:`~taichi.lang.matrix.Matrix` instance filled with ones.
 766
 767        """
 768        from taichi.lang import matrix_ops  # pylint: disable=C0415
 769
 770        if m is None:
 771            return matrix_ops._filled_vector(n, dt, 1)
 772        return matrix_ops._filled_matrix(n, m, dt, 1)
 773
 774    @staticmethod
 775    @taichi_scope
 776    def unit(n, i, dt=None):
 777        """Constructs a n-D vector with the `i`-th entry being equal to one and
 778        the remaining entries are all zeros.
 779
 780        Args:
 781            n (int): The length of the vector.
 782            i (int): The index of the entry that will be filled with one.
 783            dt (:mod:`~taichi.types.primitive_types`, optional): The desired data type.
 784
 785        Returns:
 786            :class:`~taichi.Matrix`: The returned vector.
 787
 788        Example::
 789
 790            >>> A = ti.Matrix.unit(3, 1)
 791            >>> A
 792            [0, 1, 0]
 793        """
 794        from taichi.lang import matrix_ops  # pylint: disable=C0415
 795
 796        if dt is None:
 797            dt = int
 798        assert 0 <= i < n
 799        return matrix_ops._unit_vector(n, i, dt)
 800
 801    @staticmethod
 802    @taichi_scope
 803    def identity(dt, n):
 804        """Constructs an identity Matrix with shape (n, n).
 805
 806        Args:
 807            dt (DataType): The desired data type.
 808            n (int): The number of rows/columns.
 809
 810        Returns:
 811            :class:`~taichi.Matrix`: An `n x n` identity matrix.
 812        """
 813        from taichi.lang import matrix_ops  # pylint: disable=C0415
 814
 815        return matrix_ops._identity_matrix(n, dt)
 816
 817    @classmethod
 818    @python_scope
 819    def field(
 820        cls,
 821        n,
 822        m,
 823        dtype,
 824        shape=None,
 825        order=None,
 826        name="",
 827        offset=None,
 828        needs_grad=False,
 829        needs_dual=False,
 830        layout=Layout.AOS,
 831        ndim=None,
 832    ):
 833        """Construct a data container to hold all elements of the Matrix.
 834
 835        Args:
 836            n (int): The desired number of rows of the Matrix.
 837            m (int): The desired number of columns of the Matrix.
 838            dtype (DataType, optional): The desired data type of the Matrix.
 839            shape (Union[int, tuple of int], optional): The desired shape of the Matrix.
 840            order (str, optional): order of the shape laid out in memory.
 841            name (string, optional): The custom name of the field.
 842            offset (Union[int, tuple of int], optional): The coordinate offset
 843                of all elements in a field.
 844            needs_grad (bool, optional): Whether the Matrix need grad field (reverse mode autodiff).
 845            needs_dual (bool, optional): Whether the Matrix need dual field (forward mode autodiff).
 846            layout (Layout, optional): The field layout, either Array Of
 847                Structure (AOS) or Structure Of Array (SOA).
 848
 849        Returns:
 850            :class:`~taichi.Matrix`: A matrix.
 851        """
 852        entries = []
 853        element_dim = ndim if ndim is not None else 2
 854        if isinstance(dtype, (list, tuple, np.ndarray)):
 855            # set different dtype for each element in Matrix
 856            # see #2135
 857            if m == 1:
 858                assert (
 859                    len(np.shape(dtype)) == 1 and len(dtype) == n
 860                ), f"Please set correct dtype list for Vector. The shape of dtype list should be ({n}, ) instead of {np.shape(dtype)}"
 861                for i in range(n):
 862                    entries.append(
 863                        impl.create_field_member(
 864                            dtype[i],
 865                            name=name,
 866                            needs_grad=needs_grad,
 867                            needs_dual=needs_dual,
 868                        )
 869                    )
 870            else:
 871                assert (
 872                    len(np.shape(dtype)) == 2 and len(dtype) == n and len(dtype[0]) == m
 873                ), f"Please set correct dtype list for Matrix. The shape of dtype list should be ({n}, {m}) instead of {np.shape(dtype)}"
 874                for i in range(n):
 875                    for j in range(m):
 876                        entries.append(
 877                            impl.create_field_member(
 878                                dtype[i][j],
 879                                name=name,
 880                                needs_grad=needs_grad,
 881                                needs_dual=needs_dual,
 882                            )
 883                        )
 884        else:
 885            for _ in range(n * m):
 886                entries.append(impl.create_field_member(dtype, name=name, needs_grad=needs_grad, needs_dual=needs_dual))
 887        entries, entries_grad, entries_dual = zip(*entries)
 888
 889        entries = MatrixField(entries, n, m, element_dim)
 890        if all(entries_grad):
 891            entries_grad = MatrixField(entries_grad, n, m, element_dim)
 892            entries._set_grad(entries_grad)
 893        if all(entries_dual):
 894            entries_dual = MatrixField(entries_dual, n, m, element_dim)
 895            entries._set_dual(entries_dual)
 896
 897        impl.get_runtime().matrix_fields.append(entries)
 898
 899        if shape is None:
 900            if offset is not None:
 901                raise TaichiSyntaxError("shape cannot be None when offset is set")
 902            if order is not None:
 903                raise TaichiSyntaxError("shape cannot be None when order is set")
 904        else:
 905            if isinstance(shape, numbers.Number):
 906                shape = (shape,)
 907            if isinstance(offset, numbers.Number):
 908                offset = (offset,)
 909            dim = len(shape)
 910            if offset is not None and dim != len(offset):
 911                raise TaichiSyntaxError(
 912                    f"The dimensionality of shape and offset must be the same ({dim} != {len(offset)})"
 913                )
 914            axis_seq = []
 915            shape_seq = []
 916            if order is not None:
 917                if dim != len(order):
 918                    raise TaichiSyntaxError(
 919                        f"The dimensionality of shape and order must be the same ({dim} != {len(order)})"
 920                    )
 921                if dim != len(set(order)):
 922                    raise TaichiSyntaxError("The axes in order must be different")
 923                for ch in order:
 924                    axis = ord(ch) - ord("i")
 925                    if axis < 0 or axis >= dim:
 926                        raise TaichiSyntaxError(f"Invalid axis {ch}")
 927                    axis_seq.append(axis)
 928                    shape_seq.append(shape[axis])
 929            else:
 930                axis_seq = list(range(dim))
 931                shape_seq = list(shape)
 932            same_level = order is None
 933            if layout == Layout.SOA:
 934                for e in entries._get_field_members():
 935                    impl._create_snode(axis_seq, shape_seq, same_level).place(ScalarField(e), offset=offset)
 936                if needs_grad:
 937                    for e in entries_grad._get_field_members():
 938                        impl._create_snode(axis_seq, shape_seq, same_level).place(ScalarField(e), offset=offset)
 939                if needs_dual:
 940                    for e in entries_dual._get_field_members():
 941                        impl._create_snode(axis_seq, shape_seq, same_level).place(ScalarField(e), offset=offset)
 942            else:
 943                impl._create_snode(axis_seq, shape_seq, same_level).place(entries, offset=offset)
 944                if needs_grad:
 945                    impl._create_snode(axis_seq, shape_seq, same_level).place(entries_grad, offset=offset)
 946                if needs_dual:
 947                    impl._create_snode(axis_seq, shape_seq, same_level).place(entries_dual, offset=offset)
 948        return entries
 949
 950    @classmethod
 951    @python_scope
 952    def ndarray(cls, n, m, dtype, shape):
 953        """Defines a Taichi ndarray with matrix elements.
 954        This function must be called in Python scope, and after `ti.init` is called.
 955
 956        Args:
 957            n (int): Number of rows of the matrix.
 958            m (int): Number of columns of the matrix.
 959            dtype (DataType): Data type of each value.
 960            shape (Union[int, tuple[int]]): Shape of the ndarray.
 961
 962        Example::
 963
 964            The code below shows how a Taichi ndarray with matrix elements \
 965            can be declared and defined::
 966
 967                >>> x = ti.Matrix.ndarray(4, 5, ti.f32, shape=(16, 8))
 968        """
 969        if isinstance(shape, numbers.Number):
 970            shape = (shape,)
 971        return MatrixNdarray(n, m, dtype, shape)
 972
 973    @staticmethod
 974    def rows(rows):
 975        """Constructs a matrix by concatenating a list of
 976        vectors/lists row by row. Must be called in Taichi scope.
 977
 978        Args:
 979            rows (List): A list of Vector (1-D Matrix) or a list of list.
 980
 981        Returns:
 982            :class:`~taichi.Matrix`: A matrix.
 983
 984        Example::
 985
 986            >>> @ti.kernel
 987            >>> def test():
 988            >>>     v1 = ti.Vector([1, 2, 3])
 989            >>>     v2 = ti.Vector([4, 5, 6])
 990            >>>     m = ti.Matrix.rows([v1, v2])
 991            >>>     print(m)
 992            >>>
 993            >>> test()
 994            [[1, 2, 3], [4, 5, 6]]
 995        """
 996        from taichi.lang import matrix_ops  # pylint: disable=C0415
 997
 998        return matrix_ops.rows(rows)
 999
1000    @staticmethod
1001    def cols(cols):
1002        """Constructs a Matrix instance by concatenating Vectors/lists column by column.
1003
1004        Args:
1005            cols (List): A list of Vector (1-D Matrix) or a list of list.
1006
1007        Returns:
1008            :class:`~taichi.Matrix`: A matrix.
1009
1010        Example::
1011
1012            >>> @ti.kernel
1013            >>> def test():
1014            >>>     v1 = ti.Vector([1, 2, 3])
1015            >>>     v2 = ti.Vector([4, 5, 6])
1016            >>>     m = ti.Matrix.cols([v1, v2])
1017            >>>     print(m)
1018            >>>
1019            >>> test()
1020            [[1, 4], [2, 5], [3, 6]]
1021        """
1022        from taichi.lang import matrix_ops  # pylint: disable=C0415
1023
1024        return matrix_ops.cols(cols)
1025
1026    def __hash__(self):
1027        # TODO: refactor KernelTemplateMapper
1028        # If not, we get `unhashable type: Matrix` when
1029        # using matrices as template arguments.
1030        return id(self)
1031
1032    def dot(self, other):
1033        """Performs the dot product of two vectors.
1034
1035        To call this method, both multiplicatives must be vectors.
1036
1037        Args:
1038            other (:class:`~taichi.Matrix`): The input Vector.
1039
1040        Returns:
1041            DataType: The dot product result (scalar) of the two Vectors.
1042
1043        Example::
1044
1045            >>> v1 = ti.Vector([1, 2, 3])
1046            >>> v2 = ti.Vector([3, 4, 5])
1047            >>> v1.dot(v2)
1048            26
1049        """
1050        from taichi.lang import matrix_ops  # pylint: disable=C0415
1051
1052        return matrix_ops.dot(self, other)
1053
1054    def cross(self, other):
1055        """Performs the cross product with the input vector (1-D Matrix).
1056
1057        Both two vectors must have the same dimension <= 3.
1058
1059        For two 2d vectors (x1, y1) and (x2, y2), the return value is the
1060        scalar `x1*y2 - x2*y1`.
1061
1062        For two 3d vectors `v` and `w`, the return value is the 3d vector
1063        `v x w`.
1064
1065        Args:
1066            other (:class:`~taichi.Matrix`): The input Vector.
1067
1068        Returns:
1069            :class:`~taichi.Matrix`: The cross product of the two Vectors.
1070        """
1071        from taichi.lang import matrix_ops  # pylint: disable=C0415
1072
1073        return matrix_ops.cross(self, other)
1074
1075    def outer_product(self, other):
1076        """Performs the outer product with the input Vector (1-D Matrix).
1077
1078        The outer_product of two vectors `v = (x1, x2, ..., xn)`,
1079        `w = (y1, y2, ..., yn)` is a `n` times `n` square matrix, and its `(i, j)`
1080        entry is equal to `xi*yj`.
1081
1082        Args:
1083            other (:class:`~taichi.Matrix`): The input Vector.
1084
1085        Returns:
1086            :class:`~taichi.Matrix`: The outer product of the two Vectors.
1087        """
1088        from taichi.lang import matrix_ops  # pylint: disable=C0415
1089
1090        return matrix_ops.outer_product(self, other)